import numpy as np
from scipy.linalg import svd
import json

class measureJson:
    def __init__(self):
        self.jsonName = "measurePoint.json"  #json file
        with open(self.jsonName, encoding='utf-8') as f:    #以读方式打开measurePoint.josn文件
            self.measurepoint = json.load(f)
        
    def getAllMeasurepoint(self, measureSet=None):
        if (measureSet == None):
            measureSet = self.measurepoint.keys()
            
        ret = [[self.measurepoint[i]['Coordinates']] for i in measureSet]
        return ret
 
class PointMatcher2:
    def __init__(self, set_A, set_B):
        self.set_A = set_A
        self.set_B = set_B

    def match_points(self):
        # select the corresponding points and weights from each set
        matched_A = self.set_A
        matched_B = self.set_B

        # calculate the weighted centroid of each set of points for each coordinate
        weighted_centroid_A = np.array([np.average(matched_A[:, i]) for i in range(3)])
        weighted_centroid_B = np.array([np.average(matched_B[:, i]) for i in range(3)])

        # normalize the weights to sum to 1
        # matched_weight_A = (matched_x_weight_A + matched_y_weight_A + matched_z_weight_A) / (matched_x_weight_A.sum() + 
        #                                                                                         matched_y_weight_A.sum() + matched_z_weight_A.sum())
        # matched_weight_B = (matched_x_weight_B + matched_y_weight_B + matched_z_weight_B) / (matched_x_weight_B.sum() + 
        #                                                                                         matched_y_weight_B.sum() + matched_z_weight_B.sum())

        # shift the points to the weighted centroid for each coordinate
        shifted_matched_A = matched_A - weighted_centroid_A
        shifted_matched_B = matched_B - weighted_centroid_B

        # calculate the weighted dot product for each coordinate
        # weighted_dot_products = []
        # for i in range(3):

        #     H = np.dot((shifted_matched_A[:, i] * matched_weight_A).T, shifted_matched_B[:, i] * matched_weight_B)
        #     weighted_dot_products.append(H)

        # # combine the weighted dot products for each coordinate into a single matrix
        # H = np.vstack(weighted_dot_products)
        H = np.dot(shifted_matched_A.T,shifted_matched_B)
        # perform singular value decomposition (SVD)
        U, Sig, V = np.linalg.svd(H)
        R=np.dot(V.T,U.T)

        if(np.linalg.det(R)<0):
            V[2,:]*=-1
            R = np.dot(V.T,U.T)
        
        # calculate the residual (minimum distance)
        transformed_A = np.dot((self.set_A - weighted_centroid_A), R) + weighted_centroid_B
        dist_matrix = np.linalg.norm(transformed_A - self.set_B, axis=1)

        min_index = np.argmin(dist_matrix)
        residual = dist_matrix[min_index]
        # eigvalue,vec=np.linalg.eig(R)

        return R, -np.matmul(R,weighted_centroid_A)+weighted_centroid_B, dist_matrix, residual 

class PointMatcher:
    def __init__(self, set_A, set_B, correspondences, x_weight_A=None, y_weight_A=None, z_weight_A=None, 
                 x_weight_B=None, y_weight_B=None, z_weight_B=None):
        self.set_A = set_A
        self.set_B = set_B
        self.correspondences = correspondences
        if x_weight_A is None:
            self.x_weight_A = np.ones(set_A.shape[0])
        else:
            self.x_weight_A = x_weight_A
        if y_weight_A is None:
            self.y_weight_A = np.ones(set_A.shape[0])
        else:
            self.y_weight_A = y_weight_A
        if z_weight_A is None:
            self.z_weight_A = np.ones(set_A.shape[0])
        else:
            self.z_weight_A = z_weight_A
        if x_weight_B is None:
            self.x_weight_B = np.ones(set_B.shape[0])
        else:
            self.x_weight_B = x_weight_B
        if y_weight_B is None:
            self.y_weight_B = np.ones(set_B.shape[0])
        else:
            self.y_weight_B = y_weight_B
        if z_weight_B is None:
            self.z_weight_B = np.ones(set_B.shape[0])
        else:
            self.z_weight_B = z_weight_B

    def match_points(self):
        # select the corresponding points and weights from each set
        matched_A = self.set_A[self.correspondences[:, 0]]
        matched_B = self.set_B[self.correspondences[:, 1]]
        matched_x_weight_A = self.x_weight_A[self.correspondences[:, 0]]
        matched_y_weight_A = self.y_weight_A[self.correspondences[:, 0]]
        matched_z_weight_A = self.z_weight_A[self.correspondences[:, 0]]
        matched_x_weight_B = self.x_weight_B[self.correspondences[:, 1]]
        matched_y_weight_B = self.y_weight_B[self.correspondences[:, 1]]
        matched_z_weight_B = self.z_weight_B[self.correspondences[:, 1]]

        # calculate the weighted centroid of each set of points for each coordinate
        weighted_centroid_A = np.array([np.average(matched_A[:, i], 
                                                    weights=([matched_x_weight_A, matched_y_weight_A, matched_z_weight_A])[i]) for i in range(3)])
        weighted_centroid_B = np.array([np.average(matched_B[:, i], 
                                                    weights=([matched_x_weight_B, matched_y_weight_B, matched_z_weight_B])[i]) for i in range(3)])

        # normalize the weights to sum to 1
        matched_weight_A = (matched_x_weight_A + matched_y_weight_A + matched_z_weight_A) / (matched_x_weight_A.sum() + 
                                                                                                matched_y_weight_A.sum() + matched_z_weight_A.sum())
        matched_weight_B = (matched_x_weight_B + matched_y_weight_B + matched_z_weight_B) / (matched_x_weight_B.sum() + 
                                                                                                matched_y_weight_B.sum() + matched_z_weight_B.sum())

        # shift the points to the weighted centroid for each coordinate
        shifted_matched_A = matched_A - weighted_centroid_A
        shifted_matched_B = matched_B - weighted_centroid_B

        # calculate the weighted dot product for each coordinate
        # weighted_dot_products = []
        # for i in range(3):

        #     H = np.dot((shifted_matched_A[:, i] * matched_weight_A).T, shifted_matched_B[:, i] * matched_weight_B)
        #     weighted_dot_products.append(H)

        # # combine the weighted dot products for each coordinate into a single matrix
        # H = np.vstack(weighted_dot_products)
        H = np.dot(shifted_matched_A.T,shifted_matched_B)
        # perform singular value decomposition (SVD)
        U, Sig, V = np.linalg.svd(H)
        R=np.dot(V.T,U.T)
        
        # calculate the residual (minimum distance)
        transformed_A = np.dot((self.set_A - weighted_centroid_A), R) + weighted_centroid_B
        dist_matrix = np.linalg.norm(transformed_A - self.set_B, axis=1)
        min_index = np.argmin(dist_matrix)
        residual = dist_matrix[min_index]
        eigvalue,vec=np.linalg.eig(R)
        return R, residual,np.matmul(R,weighted_centroid_A)-weighted_centroid_B,dist_matrix

if __name__ == '__main__':
    #测试
    # set_A = np.array([[1,0,1],[2,0,0],[1,1,0],[1,0,0],[2,1,1]])
    # set_B = np.array([[0,0,1],[1,0,0],[0,1,0],[0,0,0],[1,1,1]])

    #upperleft
    # set_A = np.array([[ 711.37932874,  512.02637438, -237.66942599],[ 678.54824422, -488.17487801, -265.7380096 ],[687.99912533,  19.42111162, -21.74868948],[ 696.61067122,  556.81077358, -264.10360883]])
    # set_B = np.array([[-500,0,0],[500,0,0],[0,0,-237.56],[-530,0,30]])

    #upperright
    # set_A = np.array([[-883.34219342,  566.48457542, -296.16468119],[-917.78188372, -448.82105755, -331.69311594],[-409.81100072,  549.79416565, -285.34232296],[-901.28668961,   68.70775223,  -78.85635066],[-374.40682763,  622.22833046, -816.15589621]])
    # set_B = np.array([[-500,0,0],[500,0,0],[-500,472,0],[-15,0,237.56],[-530,502,530]])

    #platform
    set_A = np.array([[-928.79809666,  254.88391323, -844.32392457],[-940.14443057,  -77.31542587, -858.86392263],[-640.09953787,  410.44153849, -828.48362489],[-925.34832494,  257.68345226, -859.31609007],[-939.29147275,  -76.13583152, -873.7972325 ]])
    set_B = np.array([[286,-165,0],[286,165.12,0],[0,-330.24,0],[286,-165.12,-15],[286,165.12,-15]])

    item = "upperright"
    numPoints = 10
    fitSet = [item + str(i) for i in range(numPoints)]
    # 定义对应的 correspondences

    correspondences = np.zeros([numPoints, 2]).astype(int)
    correspondences[:, 0] = range(numPoints)
    correspondences[:, 1] = range(numPoints)

    measurePoint = measureJson()
    set_A = measurePoint.getAllMeasurepoint(fitSet)

    ##set_B
    set_B = np.loadtxt(item+".csv")


    # 实例化 PointMatcher 类的对象
    point_matcher = PointMatcher(set_A, set_B, correspondences)
    # print(point_matcher.getAllMeasurepoint())

    # 调用 match_points() 方法
    R, residual,vec,err = point_matcher.match_points()

    print("Rotation matrix:\n", R)
    print("Residual:", residual)
    print("vec: ", vec)
    print("err:",err)
    #error view
    for i in range(numPoints):
        A = set_A[i,:]
        B = set_B[i,:]
        transformed_A = np.matmul(R,A) - vec
        err = transformed_A - B
        print(np.linalg.norm(err))

